import torch.nn as nn
from . import block as B
from . import architecture
import torch

class SRNet(nn.Module):
    def __init__(self, in_nc=3, nf=16, num_modules=3, out_nc=3, upscale=4,k=2):
        super(SRNet, self).__init__()
        self.crfn_low = architecture.CRFN(in_nc=in_nc,out_nc=out_nc,nf=nf,num_modules=num_modules,upscale= upscale,k=k)
        self.crfn_high = architecture.CRFN(in_nc=in_nc,out_nc=out_nc,nf=nf,num_modules=num_modules,upscale= upscale,k=k)

    def load_weight(self, low_path, high_path):
        low_checkpoint = torch.load(low_path)
        high_checkpoint = torch.load(high_path)
        self.crfn_low.load_state_dict(low_checkpoint)
        self.crfn_high.load_state_dict(high_checkpoint)


    def forward(self, low_freq, high_freq):
        sr_low = self.crfn_low(low_freq)
        sr_high = self.crfn_high(high_freq)
        sr_result = sr_low + sr_high
        return sr_result
